import os
import json


def get_rewards(trajectories_save_path):
    best_child_reward = []
    # best_child_em = []
    for file in os.listdir(trajectories_save_path):
        if not file.endswith('json'):
            continue
        with open(os.path.join(trajectories_save_path, file)) as f:
            result=json.load(f)
        # best_reward.append(result['best reward'])
        # best_em.append(result['best em'] if result['best em']  is not None else 0)
        best_child_reward.append(result['best child reward'])
        # best_child_em.append(result['best child em'] if result['best child em']  is not None else 0)
    return best_child_reward


# trajectories_0_save_path = "hotpot/trajectories_Phi3_1T/trajectories_test_Phi-3_simple_1iterations"
# trajectories_1_save_path = "hotpot/trajectories_Phi3_1T/trajectories-0_test_Phi-3_simple_1iterations"
# trajectories_2_save_path = "hotpot/trajectories_Phi3_1T/trajectories-1_test_Phi-3_simple_1iterations"
# trajectories_3_save_path = "hotpot/trajectories_Phi3_1T/trajectories-2_test_Phi-3_simple_1iterations"
# trajectories_4_save_path = "hotpot/trajectories_Phi3_1T/trajectories-3_test_Phi-3_simple_1iterations"


trajectories_0_save_path = "hotpot/trajectories-05T-0_test_Phi-3_simple_1iterations"
trajectories_1_save_path = "hotpot/trajectories-05T-1_test_Phi-3_simple_1iterations"
trajectories_2_save_path = "hotpot/trajectories-05T-2_test_Phi-3_simple_1iterations"
trajectories_3_save_path = "hotpot/trajectories-05T-3_test_Phi-3_simple_1iterations"
trajectories_4_save_path = "hotpot/trajectories-05T-4_test_Phi-3_simple_1iterations"



best_child_reward_0 = get_rewards(trajectories_0_save_path)
best_child_reward_1 = get_rewards(trajectories_1_save_path)
best_child_reward_2 = get_rewards(trajectories_2_save_path)
best_child_reward_3 = get_rewards(trajectories_3_save_path)
best_child_reward_4 = get_rewards(trajectories_4_save_path)

BoN_list = [max(values) for values in zip(best_child_reward_0, best_child_reward_1, best_child_reward_2,
                                           best_child_reward_3, best_child_reward_4)]

print("Sample number: ", len(BoN_list))
# print("average best reward: ", sum(best_reward)/len(best_reward))
# print("average best em: ", sum(best_em)/len(best_reward))
print("average best child reward: ", sum(BoN_list)/len(BoN_list))
# print("average best child em: ", sum(best_child_em)/len(best_child_em))